# -*- coding: utf-8 -*-
import os
import threading

from sentence_transformers import SentenceTransformer
from typing import List
from loguru import logger as log

model_dict = {}
rlock = threading.RLock()

current_dir = os.path.dirname(__file__)
model_path = os.path.join(current_dir, 'model/bge-base')


def get_embedding(sentences: List[str], model_name: str = 'model/bge-base', convert_to_numpy=True):
    """
        embedding wrapper
    """
    with rlock:
        cur_name = threading.currentThread().getName()
        log.info("get_embedding current_thread={}", cur_name)
        local_cache_model = model_dict.get(model_name, None)
        if local_cache_model:
            return local_cache_model.encode(sentences, normalize_embeddings=True, convert_to_numpy=convert_to_numpy)
        else:
            # 加载本地模型
            model = SentenceTransformer.load(model_path)
            model_dict[model_name] = model
            return model.encode(sentences, normalize_embeddings=True, convert_to_numpy=convert_to_numpy)


if __name__ == '__main__':
    sentences_1 = ["""临湖社区玉屏南路商业街区综合治理工程项目设计的招标条件是什么？
       """]
    # bge-base similarity=[[0.85091794]]
    # sentences_2 = ["""
    #         "1. 招标条件
    #
    #  1.1 项目名称：临湖社区玉屏南路商业街区综合治理工程项目设计
    #
    #  1.2 项目审批、核准或备案机关名称：合肥经济技术开发区经济发展局
    #
    #  1.3 批文名称及编号：《关于临湖社区玉屏南路商业街区综合治理工程项目立项的批复》，合经投〔2024〕41号
    #
    #  1.4 招标人：合肥经济技术开发区重点工程建设管理中心
    #
    #  1.5 项目业主：合肥经济技术开发区重点工程建设管理中心
    #
    #  1.6 资金来源：政府投资
    #
    #  1.7 项目出资比例：100%
    #
    #  1.8 资金落实情况: 已落实
    #
    # "
    #      """]
    # bge-base similarity=[[0.8576449]]
    sentences_2 = ["""
          1.1 项目名称：临湖社区玉屏南路商业街区综合治理工程项目设计 \n        \n   
         """]
    # bge-base similarity=[[0.89342594]]
    # sentences_2 = ["""
    #        "2.1 招标项目名称：临湖社区玉屏南路商业街区综合治理工程项目设计 \n        \n
    #      """]
    # embeddings_1 = get_embedding(sentences_1, 'model/bge-base', )
    # print(embeddings_1)
    # embeddings_2 = get_embedding(sentences_2, 'model/bge-base', )
    # similarity = embeddings_1 @ embeddings_2.T
    # print(f'bge-base similarity={similarity}')
    embeddings_1 = get_embedding(sentences_1, 'model/bge-base')
    embeddings_2 = get_embedding(sentences_2, 'model/bge-base')
    similarity = embeddings_1 @ embeddings_2.T
    print(f'bge-m3 similarity={similarity}')
